import pandas as pd
import os
import shutil


def copy_files(src, dst, name_list):
    for name in name_list:
        src_path = os.path.join(src, name)
        dst_path = os.path.join(dst, name)
        shutil.copy(src_path, dst_path)


def get_images_info(path, threshold=200):
    dataset = pd.read_csv(path)
    count = dataset.category.value_counts()
    print(count)
    count = count[count.values >= threshold]
    print(count)
    data = dataset[dataset.category.isin(count.index)]
    print(data)
    return data


def copy_files(src, dst, name_list):
    for name in name_list:
        src_path = os.path.join(src, name)
        dst_path = os.path.join(dst, name)
        shutil.copy(src_path, dst_path)


original_dataset_dir = '..\\archive\\images'
base_dir = 'dataset\\traffic_sign'
train_dir = os.path.join(base_dir, 'train')
os.mkdir(train_dir)
validation_dir = os.path.join(base_dir, 'validation')
os.mkdir(validation_dir)


data = get_images_info('..\\archive\\annotations.csv')
names = data.drop_duplicates(['category'])
print(names)
for name in names['category']:
    os.mkdir(os.path.join(train_dir, str(name)))
    os.mkdir(os.path.join(validation_dir, str(name)))

    images = data[data.category == name]
    copy_files(original_dataset_dir, os.path.join(train_dir, str(name)), images['file_name'][:160])
    copy_files(original_dataset_dir, os.path.join(validation_dir, str(name)), images['file_name'][160:200])
